using DataFrames
using QuadGK
using Distributions
using JLD
using Optim
using ForwardDiff
using CSV
using LinearAlgebra
using SpecialFunctions
using Random
using DelimitedFiles
using SparseArrays

#clearconsole()

#-------------------------------------------------------------
# include Functions
#-------------------------------------------------------------
#cd("$(pwd())/Dropbox/Heterogeneity/Software/KS_Simulation/")
readDir = "$(pwd())/Functions/"
include(readDir *"logSpline_Procedures.jl");
include(readDir *"fKF.jl");
include(readDir *"VAR_MoreLags_Procedures.jl");
include(readDir *"Loaddata.jl");


#-------------------------------------------------------------
# choose specification files
#-------------------------------------------------------------
nfVARSpec  = "10tc"
#nSSSpec    = "1" # K=10
#nSSSpec    = "2" # K=22
nSSSpec    = "3" # K=10, lambda2=lambda3=1.0
nSSMCMCSpec= "1"
specDir    = "$(pwd())/SpecFiles/"
include(specDir * "/fVARspec" * nfVARSpec * ".jl")
include(specDir * "/SSspec" * nSSSpec * ".jl")
include(specDir * "/SSMCMCspec" * nSSMCMCSpec * ".jl")

#-------------------------------------------------------------
# load aggregate data
#-------------------------------------------------------------
dataDir  = "$(pwd())/data/"
agg_data, period_agg, ~ = loadaggdata(SampleStart,SampleEnd)
n_agg    = size(agg_data)[2]
Tend     = size(agg_data)[1]

sNameLoadDir = "fVAR" * nfVARSpec
loaddir      = "$(pwd())/results/" * sNameLoadDir *"/";
sNameFile    = "K" * string(K) * "_fVAR" * nfVARSpec

PhatDensCoef_factor, MDD_GoF, VinvLam_all, period_Dens_ind, ~, ~, ~  = loaddensdata(SampleStart,SampleEnd,K,nfVARSpec)
n_cross = size(PhatDensCoef_factor)[2]

#-------------------------------------------------------------
# Adjust sample
#-------------------------------------------------------------

# Arrange data for VAR analysis
# We start by dropping the initial Tdrop-nlags observations
agg_data_adj = agg_data[Tdrop-nlags+1:end,:]
PhatDensCoef_factor_adj = PhatDensCoef_factor[Tdrop-nlags+1:end,:]
# We now have T*=(T-T0)+nlags observations and call the first observation t=1
# Observations t=1,...,nlags are used to initialize lags
# and observations t=nlags+1,...,T* are the estimation sample
# Note that by holding Tdrop fixed and changing nlags<=Tdrop we can change
# the nunber of lags holding the estimation sample fixed.
# Estimation sample has T*-nlags=T-T0 observations

# Define some constants
n       = n_agg+n_cross
p       = nlags

# Combine aggregate and cross sectional observations and create YY XX for VAR estimation
data_adj = [ agg_data_adj PhatDensCoef_factor_adj ]
YY, XX, PHIols, Sols, SIGMAols = data_to_YY(data_adj, nlags, nlags)
# FS: having nlags twice as input arguments is not necessary. We can simplify
# but would have to adjust other programs as well.

# transform data into companion form, used for filtering/smoothing
# _adj data: t=nlags to t=T*
comp_data  = compdata_to_WW(agg_data_adj, PhatDensCoef_factor_adj, p)
(TmT0p1,~) = size(comp_data)

#-------------------------------------------------------------
# Collect measurement error variances
#-------------------------------------------------------------

VinvLam_all_adj = VinvLam_all[:,:,Tdrop-nlags+1:end]
H_t_adj = zeros(size(VinvLam_all_adj))
for tt = 1:(TmT0p1+p-1)
    H_t_adj[:,:,tt] = inv(Symmetric(VinvLam_all[:,:,tt]))
end

#-------------------------------------------------------------
# Generate prior
#-------------------------------------------------------------

igammadf = size(YY)[2]+sigdf
prior    = prior_ACP_stru(YY, SIGMAols, n_agg, igammadf, nlags, lambda)

#-------------------------------------------------------------
# Run Posterior Sampler
#-------------------------------------------------------------

println("")
println("Bayesian estimation of State-Space Model: Gibbs Sampling... ")
println("")

# Initialize matrices for collecting draws from posterior
store_Bpdraw             = zeros(nsim,n*p,n)
store_Sigpdraw           = zeros(nsim,n,n)
store_Sigtrpdraw         = zeros(nsim,n,n)
store_alphadraw_filtered = zeros(nsim,TmT0p1+(p-1),n_cross)
store_alphadraw_smoothed = zeros(nsim,TmT0p1+(p-1),n_cross)

# Initialization
let counter = 0
    data_adj_j = [ agg_data_adj PhatDensCoef_factor_adj ]
    time_init_loop = time_ns();

    # Main Loop
    for jj = 1:nsim

        # Draw parameters given states
        YY, XX, PHIhat, Shat, SIGMAhat = data_to_YY(data_adj_j, nlags, nlags)
        alp_j, beta_j, Sig_j           = sample_ThetaSig(YY,XX,PHIhat,Shat,SIGMAhat,nlags,prior,1) # last input: nsim=1
        Btilde, Sigtilde               = getReducedForm(alp_j,beta_j,Sig_j)
        Sigtilde_j                     = Symmetric(Sigtilde[1,:,:])
        Sigtildetr_j                   = cholesky(Sigtilde_j)
        Btilde_j                       = reshape(Btilde[1,:], n*p,n)
        store_Bpdraw[jj,:,:]           = Btilde_j
        store_Sigpdraw[jj,:,:]         = Sigtilde_j
        store_Sigtrpdraw[jj,:,:]       = Sigtildetr_j.L

        # set up companion form matrix
        # Btilde is np,n and has equations in columns
        Fmat_j = CompanionForm(Btilde_j)

        # Run Kalman Filter (to obtain likelihood and filtered states)
        # we run from t=nlags (initialization) to t=T*
        s_filtered_j, ~, P_filtered_j  = KFnoX_MoreLags(H_t_adj,Sigtilde_j,Fmat_j,comp_data,n_agg,p)
        # Run the simulation smoother
        # a_sim_adj, a_sim, s_sim
        a_sim_adj_j, ~, ~ = KSmoothnoX_MoreLags(s_filtered_j,P_filtered_j,Sigtilde_j,Fmat_j,comp_data,n_agg,p)
        # update data_adj with smoothed states
        data_adj_j  = [ agg_data_adj a_sim_adj_j]

        # save results from filtering and smoothing
        a_filtered_initial_lags = reverse(reshape(s_filtered_j[1,:],n_cross,p)',dims=1)
        a_filtered_j = [a_filtered_initial_lags; s_filtered_j[2:end,1:n_cross]]
        store_alphadraw_filtered[jj,:,:] = a_filtered_j
        store_alphadraw_smoothed[jj,:,:] = a_sim_adj_j
        # CSV.write(savedir * sNameFile * "_alphas_filtered"* "_" * string(jj) * ".csv", DataFrame(alphas_filtered,:auto))

        #output
        counter = counter + 1
        if counter == ncount
           println("")
           println("Draw number:  $jj")
           println("Remaining draws:  $(nsim-jj)")
           time_loop=signed(time_ns()-time_init_loop)/1000000000
           println("Elapsed time = $(time_loop)")
           counter = 0
           time_init_loop = time_ns()
        end

    end

end

store_Bpdraw             = store_Bpdraw[nburn+1:nsim,:,:];
store_Sigpdraw           = store_Sigpdraw[nburn+1:nsim,:,:];
store_Sigtrpdraw         = store_Sigtrpdraw[nburn+1:nsim,:,:];
store_alphadraw_filtered = store_alphadraw_filtered[nburn+1:nsim,:,:];
store_alphadraw_smoothed = store_alphadraw_smoothed[nburn+1:nsim,:,:];

#-------------------------------------------------------------
# Save draws
#-------------------------------------------------------------

sName    = "fVAR" * nfVARSpec * "_SS" * nSSSpec*"_MCMC" * nSSMCMCSpec
savedir  = "$(pwd())/results/" * sName *"/";
try mkdir(savedir) catch; end
save(savedir * sName* "_PostDraws.jld", "Bpdraw", store_Bpdraw, "Sigpdraw", store_Sigpdraw, "Sigtrpdraw", store_Sigtrpdraw)
save(savedir * sName* "_StateDraws.jld", "alphadraw_filtered", store_alphadraw_filtered, "alphadraw_smoothed", store_alphadraw_smoothed)

#-------------------------------------------------------------
# MISC ITEMS
#-------------------------------------------------------------

Bpmean       = dropdims(mean(store_Bpdraw,dims=1),dims=1)
Sigpmean     = dropdims(mean(store_Sigpdraw,dims=1),dims=1)
Sigtrpmean   = dropdims(mean(store_Sigtrpdraw,dims=1),dims=1)
alphapmean_filtered = dropdims(mean(store_alphadraw_filtered,dims=1),dims=1)
alphapmean_smoothed = dropdims(mean(store_alphadraw_smoothed,dims=1),dims=1)

save(savedir * sName* "_PostMeans.jld", "Bpmean", Bpmean, "Sigpmean", Sigpmean, "Sigtrpmean", Sigtrpmean)
save(savedir * sName* "_StateMeans.jld", "alphapmean_filtered", alphapmean_filtered, "alphapmean_smoothed", alphapmean_smoothed)

# save OLS estiamtes and posterior means as CSV files
CSV.write(savedir * sName * "_PHIols.csv", DataFrame(PHIols,:auto))
CSV.write(savedir * sName * "_SIGMAols.csv", DataFrame(SIGMAols,:auto))
CSV.write(savedir * sName * "_Bpmean.csv", DataFrame(Bpmean,:auto))
CSV.write(savedir * sName * "_Sigpmean.csv", DataFrame(Sigpmean,:auto))
CSV.write(savedir * sName * "_Sigtrpmean.csv", DataFrame(Sigtrpmean,:auto))
